
import os
import json
from dashscope import Generation

# JSON Schema定义
json_schema = {
    "title": "Date",
    "description": "Formated date expression",
    "type": "object",
    "properties": {
        "year": {
            "type": "integer",
            "description": "year, YYYY",
        },
        "month": {
            "type": "integer",
            "description": "month, MM",
        },
        "day": {
            "type": "integer",
            "description": "day, DD",
        },
        "era": {
            "type": "string",
            "description": "BC or AD",
        },
    },
}

# 从环境变量获取API密钥
api_key = os.getenv("DASHSCOPE_API_KEY")
if not api_key:
    api_key = "sk-2c036d80d14a443f8769e072bb09a2e1"

def extract_date_with_json_schema(query):
    """使用JSON Schema提取日期信息"""
    
    prompt = f"""提取用户输入中的日期信息，并严格按照以下JSON Schema格式返回结果：

{json.dumps(json_schema, indent=2, ensure_ascii=False)}

用户输入: {query}

请只返回符合上述Schema的JSON格式结果，不要包含其他文字。"""

    try:
        response = Generation.call(
            model='qwen-plus',
            prompt=prompt,
            api_key=api_key
        )
        
        response_text = response.output.text.strip()
        
        # 提取JSON部分
        if "```json" in response_text:
            json_start = response_text.find("```json") + 7
            json_end = response_text.find("```", json_start)
            json_text = response_text[json_start:json_end].strip()
        elif "{" in response_text and "}" in response_text:
            json_start = response_text.find("{")
            json_end = response_text.rfind("}") + 1
            json_text = response_text[json_start:json_end]
        else:
            json_text = response_text
        
        # 解析JSON
        result = json.loads(json_text)
        return result
        
    except Exception as e:
        print(f"调用模型时出错: {e}")
        return None

# 测试JSON Schema结构化输出
print("=== JSON Schema 结构化输出测试 ===")
query = "2023年四月6日天气晴..."
result = extract_date_with_json_schema(query)
if result:
    print(f"提取结果: {result}")
else:
    print("提取失败")


# 工具函数定义
def add(a: int, b: int) -> int:
    """Add two integers."""
    return a + b

def multiply(a: float, b: float) -> float:
    """Multiply two numbers."""
    return a * b

def simulate_tool_calling(query):
    """模拟工具调用功能"""
    
    # 分析查询是否需要工具调用
    if "加" in query or "add" in query.lower() or "+" in query:
        # 简单的数字提取（实际应用中需要更复杂的NLP）
        numbers = []
        for word in query.split():
            try:
                num = float(word)
                numbers.append(num)
            except ValueError:
                continue
        
        if len(numbers) >= 2:
            result = add(int(numbers[0]), int(numbers[1]))
            return {
                "tool_name": "add",
                "arguments": {"a": int(numbers[0]), "b": int(numbers[1])},
                "result": result
            }
    
    elif "倍" in query or "multiply" in query.lower() or "乘" in query:
        # 提取数字进行计算
        numbers = []
        for word in query.split():
            try:
                num = float(word)
                numbers.append(num)
            except ValueError:
                continue
        
        if len(numbers) >= 2:
            result = multiply(numbers[0], numbers[1])
            return {
                "tool_name": "multiply", 
                "arguments": {"a": numbers[0], "b": numbers[1]},
                "result": result
            }
    
    # 如果没有匹配到工具，使用AI回答
    prompt = f"""用户问题: {query}

请回答这个问题。如果涉及数学计算，请提供详细的计算过程。"""

    try:
        response = Generation.call(
            model='qwen-plus',
            prompt=prompt,
            api_key=api_key
        )
        return {
            "tool_name": "ai_response",
            "arguments": {},
            "result": response.output.text
        }
    except Exception as e:
        return {
            "tool_name": "error",
            "arguments": {},
            "result": f"调用模型时出错: {e}"
        }

# 测试工具调用功能
print("\n=== 工具调用功能测试 ===")
query = "3.5的4倍是多少?"
print(f"用户查询: {query}")

# 模拟工具调用
tool_result = simulate_tool_calling(query)
print("工具调用结果:")
print(json.dumps(tool_result, indent=4, ensure_ascii=False))

# 测试加法工具
print("\n=== 加法工具测试 ===")
add_query = "5加3等于多少"
print(f"用户查询: {add_query}")
add_result = simulate_tool_calling(add_query)
print("工具调用结果:")
print(json.dumps(add_result, indent=4, ensure_ascii=False))
